#include "space_hash_cpu.h"
#include <ctime>

namespace PhysLeo
{

template<typename T>
SpaceHashCpu<T>::SpaceHashCpu():
    num_(0),
    nx_(0),
    ny_(0),
    nz_(0),
    np_max_(0),
    ids_(nullptr),
    counter_(nullptr)
{
}

template<typename T>
SpaceHashCpu<T>::~SpaceHashCpu()
{
}

template<typename T>
void SpaceHashCpu<T>::setSpace(T resolution, glm::tvec3<T> low, glm::tvec3<T> high, const int np_max)
{
    resolution_ = resolution;
    low_ = low;
    high_ = high;
    np_max_ = np_max;

    glm::tvec3<T> num_segment = (high_ - low_) / resolution_;

    nx_ = static_cast<int>(ceil(num_segment.x));
    ny_ = static_cast<int>(ceil(num_segment.y));
    nz_ = static_cast<int>(ceil(num_segment.z));

    num_ = nx_ * ny_*nz_;
    ids_ = std::make_shared<BufferData<int>>(num_*np_max_);
    ids_->allocateCpu();
    counter_ = std::make_shared<BufferData<int>>(num_);
    counter_->allocateCpu();
}

template<typename T>
void SpaceHashCpu<T>::constructHashTable(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos)
{
    counter_->setZeroCpu();
    ids_->setZeroCpu();

    const int pos_num = ptr_pos->size();
    auto pos_ptr = ptr_pos->ptrCpu();
    auto counter_ptr = counter_->ptrCpu();
    auto ids_ptr = ids_->ptrCpu();

    for(auto i=0;i<pos_num;++i)
    {
        glm::tvec3<int> ijk = glm::floor((pos_ptr[i] - low_) / resolution_);
        int gid = ijk.x*ny_*nz_ + ijk.y*nz_ + ijk.z;
        counter_ptr[gid] = counter_ptr[gid] < np_max_ ? counter_ptr[gid] + 1 : np_max_;
        ids_ptr[gid*np_max_ + counter_ptr[gid] - 1] = i;
    }
}

template<typename T>
void SpaceHashCpu<T>::queryNeighbor(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos, std::shared_ptr<BufferData<NeighborList>> ptr_neighbor_list, T max_radius, T min_radius)
{
    constructHashTable(ptr_pos);
    ptr_neighbor_list->setZeroCpu();

    const auto start = clock();

    const int pos_num = ptr_pos->size();
    auto pos_ptr = ptr_pos->ptrCpu();
    auto counter_ptr = counter_->ptrCpu();
    auto ids_ptr = ids_->ptrCpu();
    const auto neighbor_list_ptr = ptr_neighbor_list->ptrCpu();

    for(auto i=0;i<pos_num;++i)
    {
        int ids[SPACE_HASH_MAX_NEIGHBOR];
        T distance[SPACE_HASH_MAX_NEIGHBOR];

        const glm::tvec3<int> gid3 = glm::floor((pos_ptr[i] - low_) / resolution_);

        auto counter = 0;
        for(auto dx = -1;dx<=1;++dx)
            for(auto dy=-1;dy<=1;++dy)
                for(auto dz=-1;dz<=1;++dz)
                {
                    auto gid = (gid3.x + dx)*ny_*nz_ + (gid3.y + dy)*nz_ + (gid3.z + dz);
                    if (gid < 0 || gid >= num_)continue;
                    for(auto j=0;j<counter_ptr[gid];++j)
                    {
                        int dst_index = ids_ptr[gid*np_max_ + j];
                        T distance_ = glm::length(pos_ptr[dst_index] - pos_ptr[i]);
                        if(distance_<max_radius&&distance_>min_radius)
                        {
                            if(counter< SPACE_HASH_MAX_NEIGHBOR)
                            {
                                ids[counter] = dst_index;
                                distance[counter] = distance_;
                                counter++;
                            }
                            else
                            {
                                auto max_id = 0;
                                T max_distance = distance[0];
                                for(auto k=1;k<counter;++k)
                                {
                                    if(max_distance<distance[k])
                                    {
                                        max_distance = distance[k];
                                        max_id = k;
                                    }
                                }
                                if(distance_<max_distance)
                                {
                                    ids[max_id] = dst_index;
                                    distance[max_id] = distance_;
                                }
                            }
                        }
                    }
                }
 
        neighbor_list_ptr[i].size_ = counter;
        for (auto j = 0; j < counter; ++j)
        {
            neighbor_list_ptr[i].ids_[j] = ids[j];
        }
    }

    const auto end = clock();
    std::cout << "time cost for query neighbors cpu: " << static_cast<double>(end-start)/CLOCKS_PER_SEC << std::endl;
}

template<typename T>
std::vector<glm::tvec3<T>> SpaceHashCpu<T>::bbox(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos)
{
    const int pos_num = ptr_pos->size();
    auto pos_ptr = ptr_pos->ptrCpu();

    std::vector<glm::tvec3<T>> result(2,pos_ptr[0]);
    for(auto i=1;i<pos_num;++i)
    {
        result[0] = min(result[0], pos_ptr[i]);
        result[1] = max(result[1], pos_ptr[i]);
    }
    return result;
}

template class SpaceHashCpu<float>;
template class SpaceHashCpu<double>;

}